#!/usr/bin/env python
#coding=utf-8

import json

from pprint import pprint
from aliyunsdkcore.client import AcsClient
from aliyunsdkcore.acs_exception.exceptions import ClientException
from aliyunsdkcore.acs_exception.exceptions import ServerException
from aliyunsdkcore.auth.credentials import AccessKeyCredential
from aliyunsdkcore.auth.credentials import StsTokenCredential
from aliyunsdkecs.request.v20140526.AddTagsRequest import AddTagsRequest
from aliyunsdkecs.request.v20140526.DescribeDisksRequest import DescribeDisksRequest
from core.logger import logger
import pandas as pd

from common.aliyun.aliyun_host import AliYunHost
from common.aliyun.aliyun_client import AliYunClient
from common.aliyun.aliyun_common import AliYunCommon



class AliYunDisk(object):
    def __init__(self):
        self.aliyun_host = AliYunHost()
        self.aliyun_client = AliYunClient()
        self.aliyun = self.aliyun_client.get_all_client()
        self.aliyun_common = AliYunCommon()
        


    def get_disk_info(self, account):
        page_size = 100
        page_number = 1
        disks_info = []
        while True:
            request = DescribeDisksRequest()
            request.set_accept_format('json')

            request.set_PageNumber(page_number)
            request.set_PageSize(page_size)
            try:
                response = account["client"].do_action_with_exception(request)
                response = json.loads(str(response, encoding='utf-8'))

                for disk in response["Disks"]["Disk"]:
                    tag = self.aliyun_common.get_tag(disk)
                    disk_info = {
                        "instance_id": disk["DiskId"],
                        "host_instance_id": disk["InstanceId"],
                        "name": disk["DiskName"],
                        "category": disk["Category"],
                        "size": disk["Size"],
                        "device": disk["Device"],
                        "tag": tag,
                        "instance_charge_type": disk['DiskChargeType'],
                        "status": disk["Status"],
                        "cloud_platform": "aliyun",
                        "account_name": account["name"],
                    }
                    # 增加价格信息，磁盘类型改为中文
                    self.update_instance_price(disk_info)

                    disks_info.append(disk_info)

                if page_number * page_size >= response["TotalCount"]:
                    break
                else:
                    page_number += 1

            except Exception as e:
                logger.error("get_disk_info error =>{}".format(e))
                return None

        return disks_info


    def get_all_disks_info(self):
        disks_info = []
        for account in self.aliyun:
            disk_info = self.get_disk_info(account)
            if not disk_info:
                logger.error(f"get disk info failed =>{account['name']}")
                return None
            disks_info.extend(disk_info)
            
        logger.info("get all disk success")
        return disks_info


    def update_instance_price(self, disk_info):
        """
        根据磁盘1G的价格,直接算出磁盘价格,没有调阿里云接口。
        """
        disk_price ={
            "cloud": {
                "name": "普通云盘",
                "price": 0.35
            },
            "cloud_efficiency": {
                "name": "高效云盘",
                "price": 0.35
            },
            "cloud_ssd": {
                "name": "SSD盘",
                "price": 1
            },
            "cloud_essd": {
                "name": "ESSD云盘",
                "price": 1
            },
            "cloud_auto": {
                "name": "ESSD AutoPL云盘",
                "price": 1
            }
        }
        if disk_info["category"] in disk_price:
            category_info = disk_price[disk_info["category"]]
            disk_info["category"] = category_info["name"]
            disk_info["instance_price"] = round(disk_info["size"]*category_info["price"]*0.7)
        else:
            disk_info["instance_price"] = 0
            
        return True
        


    def get_disk_info_to_excel(self):
        disks_info = self.get_all_disks_info()
        hosts_info = {
            host["id"]: host
            for host in self.aliyun_host.get_all_hosts_info()
        }

        tag_departments = self.dev_ops_api.get_all_tag_department()
        tag_infos = {
            tag_department["tag_value"]: tag_department
            for tag_department in tag_departments["data"]["items"]
        }

        disk_price ={
            "cloud": {
                "name": "普通云盘",
                "price": 0.35
            },
            "cloud_efficiency": {
                "name": "高效云盘",
                "price": 0.35
            },
            "cloud_ssd": {
                "name": "SSD盘",
                "price": 1
            },
            "cloud_essd": {
                "name": "ESSD云盘",
                "price": 1
            },
            "cloud_auto": {
                "name": "ESSD AutoPL云盘",
                "price": 1
            }
        }

        # 获取磁盘绑定的主机信息
        for disk_info in disks_info:
            if disk_info["attachment"] and disk_info["attachment"] in hosts_info:
                disk_info["host_name"] = hosts_info[disk_info["attachment"]]["host_name"]
                disk_info["host_ip"] = hosts_info[disk_info["attachment"]]["host_ip"]
            else:
                disk_info["host_name"] = ""
                disk_info["host_ip"] = ""

            if disk_info["tag"] == "p_common" or disk_info["tag"] not in tag_infos:
                disk_info["belong_department"] = "无"
                disk_info["maintain_department"] = "无"
            else:
                disk_info["belong_department"] = tag_infos[disk_info["tag"]]["belong_department"]
                disk_info["maintain_department"] = tag_infos[disk_info["tag"]]["maintain_department"]

            if disk_info["category"] in disk_price:
                disk_info["category_name"] = disk_price[disk_info["category"]]["name"]
                disk_info["price"] = disk_info["size"]*disk_price[disk_info["category"]]["price"]*0.7
            else:
                disk_info["price"] = 0

                # 生成excel文件
        excel_writer = pd.ExcelWriter("disk信息.xlsx")
        df = pd.DataFrame(disks_info)
        df.to_excel(excel_writer, sheet_name="disk", index=False)
        excel_writer.close()



    def check_all_disks_info(self):
        """
        检查所有的磁盘
        """
        disks_info = self.get_all_disks_info()

        # 检查磁盘有没绑定到服务器
        for disk in disks_info:
            if disk["status"] != "In_use":
                logger.error("{} => {} not attachment to ECS".format(disk["account"], disk["disk_name"]))
        logger.info("check disk attachment complete")

        # 检查磁盘有没打标签
        for disk in disks_info:
            if disk["tag"] == "p_common":
                logger.error("{} => {}:{} not tag".format(disk["account"], disk["id"], disk["disk_name"]))
        logger.info("check disk tag complete")
        return True



    def update_tag(self, account, disk):
        try:
            request = AddTagsRequest()
            request.set_accept_format('json')

            request.set_ResourceType("disk")
            request.set_ResourceId(disk["id"])
            request.set_Tags([
                {
                    "Key": "appname",
                    "Value": disk["tag"]
                }
            ])

            response = account["client"].do_action_with_exception(request)
            logger.info("{}=>{} OK".format(disk["disk_name"], disk["tag"]))
        except Exception as e:
            logger.error("{}=>{} error =>{}".format(disk["disk_name"], disk["tag"], e))
            return False

        return True



    def update_appname_tag(self):
        """
        更新appanme标签
        """
        hosts_info = {
            host["id"]: host
            for host in self.aliyun_host.get_all_hosts_info()
        }

        for account in self.aliyun:
            # 获取账号下所有的磁盘信息
            disk_info = self.get_disk_info(account)
            if not disk_info:
                logger.error("get disk info failed =>{}".format(account["name"]))
                continue

            # 遍历每个磁盘
            for disk in disk_info:
                #没有打标签的磁盘，打上和服务器一样的标签
                # if disk["tag"] == "p_common" and disk["status"] == "In_use":
                # 所有磁盘打上和服务器一样的标签
                if disk["status"] == "In_use":
                    host_id = disk["attachment"]

                    if host_id not in hosts_info:
                        logger.warning(f"{disk['disk_name']} not attachment ECS ")
                        continue

                    disk["tag"] = hosts_info[host_id]["tag"]
                    self.update_tag(account, disk)




def main():
    aliyun_disk = AliYunDisk()
    # aliyun_disk.update_appname_tag()
    aliyun_disk.get_all_disks_info()
    # aliyun_disk.check_all_disks_info()
    # aliyun_disk.get_disk_info_to_excel()

if __name__ == "__main__":
    main()